#!/usr/bin/env python3
"""
摄像头捕获模块
支持USB摄像头
"""

import cv2
import numpy as np
import time
import threading
import logging
from typing import Generator, Tuple, Optional
import queue

logger = logging.getLogger(__name__)

class CameraCapture:
    """摄像头捕获类"""
    
    def __init__(self, 
                 camera_id: int = 0,
                 resolution: Tuple[int, int] = (320, 240),
                 fps: int = 30,
                 use_pi_camera: bool = False):
        """
        初始化摄像头
        
        Args:
            camera_id: 摄像头ID (USB摄像头使用)
            resolution: 分辨率 (宽, 高)
            fps: 帧率
            use_pi_camera: 是否使用树莓派摄像头
        """
        self.camera_id = camera_id
        self.resolution = resolution
        self.fps = fps
        self.use_pi_camera = False
        
        self.cap = None
        self.is_running = False
        self.frame_queue = queue.Queue(maxsize=5)
        self.capture_thread = None
        
        # 性能统计
        self.frame_count = 0
        self.start_time = time.time()
        
        self._initialize_camera()
    
    def _initialize_camera(self):
        """初始化摄像头"""
        try:
            self._init_usb_camera()
            logger.info(f"USB摄像头初始化成功 (ID: {self.camera_id})")
        except Exception as e:
            logger.error(f"摄像头初始化失败: {e}")
            raise
    
    def _init_pi_camera(self):
        raise RuntimeError("不支持树莓派摄像头")
    
    def _init_usb_camera(self):
        """初始化USB摄像头"""
        self.cap = cv2.VideoCapture(self.camera_id)
        
        if not self.cap.isOpened():
            raise RuntimeError(f"无法打开摄像头 {self.camera_id}")
        
        # 设置摄像头参数
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.resolution[0])
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.resolution[1])
        self.cap.set(cv2.CAP_PROP_FPS, self.fps)
        
        # 设置缓冲区大小
        self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        
        # 验证设置
        actual_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        actual_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        actual_fps = self.cap.get(cv2.CAP_PROP_FPS)
        
        logger.info(f"USB摄像头配置: {actual_width}x{actual_height}@{actual_fps}fps")
    
    def _capture_pi_camera(self):
        raise RuntimeError("不支持树莓派摄像头")
    
    def _capture_usb_camera(self):
        """USB摄像头捕获线程"""
        while self.is_running:
            ret, frame = self.cap.read()
            
            if not ret:
                logger.warning("摄像头读取失败")
                continue
            
            # 放入队列
            if not self.frame_queue.full():
                self.frame_queue.put(frame.copy())
            
            self.frame_count += 1
            
            # 控制帧率
            time.sleep(1.0 / self.fps)
    
    def start_capture(self):
        """开始捕获"""
        if self.is_running:
            return
        
        self.is_running = True
        self.start_time = time.time()
        
        self.capture_thread = threading.Thread(target=self._capture_usb_camera)
        
        self.capture_thread.start()
        logger.info("摄像头捕获开始")
    
    def get_frame(self) -> Optional[np.ndarray]:
        """获取一帧图像"""
        try:
            return self.frame_queue.get(timeout=1.0)
        except queue.Empty:
            return None
    
    def capture_frames(self) -> Generator[np.ndarray, None, None]:
        """生成器方式获取帧"""
        self.start_capture()
        
        try:
            while self.is_running:
                frame = self.get_frame()
                if frame is not None:
                    yield frame
        finally:
            self.stop()
    
    def get_fps_stats(self) -> dict:
        """获取FPS统计"""
        elapsed_time = time.time() - self.start_time
        actual_fps = self.frame_count / elapsed_time if elapsed_time > 0 else 0
        
        return {
            'actual_fps': actual_fps,
            'target_fps': self.fps,
            'frame_count': self.frame_count,
            'elapsed_time': elapsed_time
        }
    
    def stop(self):
        """停止捕获"""
        if not self.is_running:
            return
        
        logger.info("正在停止摄像头捕获...")
        self.is_running = False
        
        if self.capture_thread and self.capture_thread.is_alive():
            self.capture_thread.join(timeout=2.0)
        
        # 释放资源
        if self.use_pi_camera and hasattr(self, 'pi_camera'):
            self.pi_camera.close()
        elif self.cap:
            self.cap.release()
        
        # 打印统计信息
        stats = self.get_fps_stats()
        logger.info(f"摄像头统计: {stats['frame_count']} 帧, "
                   f"实际FPS: {stats['actual_fps']:.1f}")
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.stop()

class CameraCalibration:
    """摄像头标定类"""
    
    def __init__(self, camera: CameraCapture):
        self.camera = camera
        self.calibration_data = None
    
    def calibrate(self, 
                  checkerboard_size: Tuple[int, int] = (9, 6),
                  num_images: int = 20) -> dict:
        """
        摄像头标定
        
        Args:
            checkerboard_size: 棋盘格内角点数量 (宽, 高)
            num_images: 标定图像数量
            
        Returns:
            标定结果字典
        """
        logger.info(f"开始摄像头标定，需要 {num_images} 张图像")
        
        # 准备标定点
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        
        # 3D点
        objp = np.zeros((checkerboard_size[0] * checkerboard_size[1], 3), np.float32)
        objp[:, :2] = np.mgrid[0:checkerboard_size[0], 0:checkerboard_size[1]].T.reshape(-1, 2)
        
        objpoints = []  # 3D点
        imgpoints = []  # 2D点
        
        captured_count = 0
        
        for frame in self.camera.capture_frames():
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 寻找棋盘格角点
            ret, corners = cv2.findChessboardCorners(gray, checkerboard_size, None)
            
            if ret:
                objpoints.append(objp)
                
                # 精确化角点
                corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
                imgpoints.append(corners2)
                
                # 绘制角点
                cv2.drawChessboardCorners(frame, checkerboard_size, corners2, ret)
                captured_count += 1
                
                logger.info(f"捕获标定图像 {captured_count}/{num_images}")
            
            # 显示图像
            cv2.imshow('Camera Calibration', frame)
            
            if cv2.waitKey(1) & 0xFF == ord('q') or captured_count >= num_images:
                break
        
        cv2.destroyAllWindows()
        
        if captured_count < 10:
            logger.error("标定图像数量不足")
            return None
        
        # 执行标定
        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(
            objpoints, imgpoints, gray.shape[::-1], None, None
        )
        
        if ret:
            self.calibration_data = {
                'camera_matrix': mtx,
                'distortion_coefficients': dist,
                'rotation_vectors': rvecs,
                'translation_vectors': tvecs,
                'calibration_error': ret
            }
            
            logger.info(f"摄像头标定完成，重投影误差: {ret:.3f}")
            return self.calibration_data
        else:
            logger.error("摄像头标定失败")
            return None
    
    def undistort_frame(self, frame: np.ndarray) -> np.ndarray:
        """去畸变"""
        if self.calibration_data is None:
            return frame
        
        mtx = self.calibration_data['camera_matrix']
        dist = self.calibration_data['distortion_coefficients']
        
        return cv2.undistort(frame, mtx, dist, None, mtx)

def test_camera():
    """测试摄像头功能"""
    logger.info("开始摄像头测试")
    
    try:
        # 创建摄像头实例
        camera = CameraCapture(resolution=(320, 240), fps=30)
        
        # 测试捕获
        frame_count = 0
        start_time = time.time()
        
        for frame in camera.capture_frames():
            frame_count += 1
            
            # 显示帧
            cv2.imshow('Camera Test', frame)
            
            # 每秒显示一次统计
            if frame_count % 30 == 0:
                elapsed = time.time() - start_time
                fps = frame_count / elapsed
                logger.info(f"测试FPS: {fps:.1f}")
            
            # 按q退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            
            # 测试60秒后自动退出
            if time.time() - start_time > 60:
                break
        
        cv2.destroyAllWindows()
        
        # 显示最终统计
        stats = camera.get_fps_stats()
        logger.info(f"测试完成: {stats}")
        
    except Exception as e:
        logger.error(f"摄像头测试失败: {e}")

if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 运行测试
    test_camera()